#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
IP + 时间窗口用户识别机制
根据腾讯云HAI成本优化方案实现用户识别和公平队列排序
"""

import time
import hashlib
from typing import Dict, Tuple, Optional
from datetime import datetime, timedelta
import logging

logger = logging.getLogger(__name__)


class UserIdentification:
    """IP + 时间窗口用户识别器"""
    
    def __init__(self, window_size_seconds: int = 3600):
        """
        初始化用户识别器
        
        Args:
            window_size_seconds: 时间窗口大小，默认1小时(3600秒)
        """
        self.window_size_seconds = window_size_seconds
        self.user_cache: Dict[str, Dict] = {}
        
    def get_time_window_id(self, timestamp: Optional[float] = None) -> int:
        """
        获取时间窗口ID
        
        Args:
            timestamp: 时间戳，如果为None则使用当前时间
            
        Returns:
            时间窗口ID
        """
        if timestamp is None:
            timestamp = time.time()
        
        # 计算时间窗口编号
        window_id = int(timestamp // self.window_size_seconds)
        return window_id
    
    def identify_user(self, client_ip: str, timestamp: Optional[float] = None) -> str:
        """
        识别用户身份
        
        Args:
            client_ip: 客户端IP地址
            timestamp: 时间戳，如果为None则使用当前时间
            
        Returns:
            用户ID，格式：user_{IP地址}_{时间窗口编号}
        """
        if timestamp is None:
            timestamp = time.time()
            
        window_id = self.get_time_window_id(timestamp)
        user_id = f"user_{client_ip}_{window_id}"
        
        # 记录用户识别信息
        if user_id not in self.user_cache:
            self.user_cache[user_id] = {
                'ip': client_ip,
                'window_id': window_id,
                'first_seen': timestamp,
                'last_seen': timestamp,
                'task_count': 0
            }
        else:
            self.user_cache[user_id]['last_seen'] = timestamp
            
        logger.info(f"用户识别: {client_ip} -> {user_id} (窗口{window_id})")
        return user_id
    
    def get_user_info(self, user_id: str) -> Optional[Dict]:
        """
        获取用户信息
        
        Args:
            user_id: 用户ID
            
        Returns:
            用户信息字典，如果用户不存在则返回None
        """
        return self.user_cache.get(user_id)
    
    def increment_user_task_count(self, user_id: str) -> int:
        """
        增加用户任务计数
        
        Args:
            user_id: 用户ID
            
        Returns:
            当前任务计数
        """
        if user_id in self.user_cache:
            self.user_cache[user_id]['task_count'] += 1
            return self.user_cache[user_id]['task_count']
        return 0
    
    def cleanup_expired_users(self, current_time: Optional[float] = None) -> int:
        """
        清理过期的用户缓存
        
        Args:
            current_time: 当前时间戳
            
        Returns:
            清理的用户数量
        """
        if current_time is None:
            current_time = time.time()
            
        # 清理超过2个时间窗口的用户
        expired_window_threshold = self.get_time_window_id(current_time) - 2
        
        expired_users = []
        for user_id, user_info in self.user_cache.items():
            if user_info['window_id'] < expired_window_threshold:
                expired_users.append(user_id)
        
        for user_id in expired_users:
            del self.user_cache[user_id]
            
        logger.info(f"清理过期用户: {len(expired_users)}个")
        return len(expired_users)
    
    def get_active_users_count(self) -> int:
        """获取活跃用户数量"""
        return len(self.user_cache)
    
    def get_user_statistics(self) -> Dict:
        """
        获取用户统计信息
        
        Returns:
            统计信息字典
        """
        if not self.user_cache:
            return {
                'total_users': 0,
                'total_tasks': 0,
                'average_tasks_per_user': 0,
                'window_distribution': {}
            }
        
        total_tasks = sum(user['task_count'] for user in self.user_cache.values())
        window_distribution = {}
        
        for user_info in self.user_cache.values():
            window_id = user_info['window_id']
            if window_id not in window_distribution:
                window_distribution[window_id] = 0
            window_distribution[window_id] += 1
        
        return {
            'total_users': len(self.user_cache),
            'total_tasks': total_tasks,
            'average_tasks_per_user': total_tasks / len(self.user_cache) if self.user_cache else 0,
            'window_distribution': window_distribution
        }


class UserFairQueue:
    """用户公平队列排序机制"""
    
    def __init__(self, user_identifier: UserIdentification):
        """
        初始化公平队列
        
        Args:
            user_identifier: 用户识别器实例
        """
        self.user_identifier = user_identifier
        self.user_queues: Dict[str, list] = {}  # 每个用户的任务队列
        self.global_queue: list = []  # 全局任务队列，按用户轮询顺序
        self.user_last_served: Dict[str, float] = {}  # 用户最后被服务的时间
        
    def add_task(self, task_id: str, client_ip: str, task_data: Dict, 
                 timestamp: Optional[float] = None) -> str:
        """
        添加任务到队列
        
        Args:
            task_id: 任务ID
            client_ip: 客户端IP
            task_data: 任务数据
            timestamp: 时间戳
            
        Returns:
            用户ID
        """
        if timestamp is None:
            timestamp = time.time()
            
        # 识别用户
        user_id = self.user_identifier.identify_user(client_ip, timestamp)
        
        # 增加用户任务计数
        self.user_identifier.increment_user_task_count(user_id)
        
        # 创建任务对象
        task = {
            'task_id': task_id,
            'user_id': user_id,
            'client_ip': client_ip,
            'task_data': task_data,
            'submit_time': timestamp,
            'status': 'queued'
        }
        
        # 添加到用户队列
        if user_id not in self.user_queues:
            self.user_queues[user_id] = []
            self.user_last_served[user_id] = 0
            
        self.user_queues[user_id].append(task)
        
        # 重新构建全局队列（公平轮询）
        self._rebuild_global_queue()
        
        logger.info(f"任务已添加到队列: {task_id} -> {user_id}, 队列长度: {len(self.global_queue)}")
        return user_id
    
    def get_next_task(self) -> Optional[Dict]:
        """
        获取下一个要处理的任务（公平轮询）
        
        Returns:
            任务字典，如果没有任务则返回None
        """
        if not self.global_queue:
            return None
            
        # 从全局队列头部取出任务
        task = self.global_queue.pop(0)
        user_id = task['user_id']
        
        # 更新用户最后被服务时间
        self.user_last_served[user_id] = time.time()
        
        # 标记任务为已获取，但不立即移除
        task['status'] = 'fetched'
        task['fetch_time'] = time.time()
        
        # 如果用户队列为空，清理相关记录
        if user_id in self.user_queues and not self.user_queues[user_id]:
            del self.user_queues[user_id]
            if user_id in self.user_last_served:
                del self.user_last_served[user_id]
        
        logger.info(f"获取任务: {task['task_id']} -> {user_id}")
        return task
    
    def _rebuild_global_queue(self):
        """重新构建全局队列，实现公平轮询"""
        # 按用户最后被服务时间排序，确保公平性
        sorted_users = sorted(
            self.user_queues.keys(),
            key=lambda uid: self.user_last_served.get(uid, 0)
        )
        
        # 重新构建全局队列
        new_global_queue = []
        max_rounds = max(len(queue) for queue in self.user_queues.values()) if self.user_queues else 0
        
        for round_num in range(max_rounds):
            for user_id in sorted_users:
                if user_id in self.user_queues and len(self.user_queues[user_id]) > round_num:
                    new_global_queue.append(self.user_queues[user_id][round_num])
        
        self.global_queue = new_global_queue
    
    def get_queue_status(self) -> Dict:
        """
        获取队列状态信息
        
        Returns:
            队列状态字典
        """
        total_tasks = sum(len(queue) for queue in self.user_queues.values())
        
        return {
            'total_tasks_in_queue': total_tasks,
            'global_queue_length': len(self.global_queue),
            'active_users': len(self.user_queues),
            'user_queue_lengths': {uid: len(queue) for uid, queue in self.user_queues.items()},
            'user_statistics': self.user_identifier.get_user_statistics()
        }
    
    def cleanup_expired_tasks(self, max_wait_time: int = 300) -> int:
        """
        清理超时任务
        
        Args:
            max_wait_time: 最大等待时间（秒）
            
        Returns:
            清理的任务数量
        """
        current_time = time.time()
        cleaned_count = 0
        
        for user_id, queue in list(self.user_queues.items()):
            expired_tasks = []
            for task in queue:
                if current_time - task['submit_time'] > max_wait_time:
                    expired_tasks.append(task)
                    cleaned_count += 1
            
            # 移除过期任务
            for task in expired_tasks:
                queue.remove(task)
            
            # 如果用户队列为空，清理相关记录
            if not queue:
                del self.user_queues[user_id]
                if user_id in self.user_last_served:
                    del self.user_last_served[user_id]
        
        # 重新构建全局队列
        self._rebuild_global_queue()
        
        if cleaned_count > 0:
            logger.warning(f"清理超时任务: {cleaned_count}个")
        
        return cleaned_count
    
    def get_task_info(self, task_id: str) -> Optional[Dict]:
        """
        获取任务信息
        
        Args:
            task_id: 任务ID
            
        Returns:
            任务信息字典，如果不存在则返回None
        """
        # 在所有用户队列中查找任务
        for user_id, queue in self.user_queues.items():
            for task in queue:
                if task['task_id'] == task_id:
                    # 计算队列位置
                    queue_position = 0
                    for i, global_task in enumerate(self.global_queue):
                        if global_task['task_id'] == task_id:
                            queue_position = i
                            break
                    
                    # 返回任务信息
                    task_info = task.copy()
                    task_info['queue_position'] = queue_position
                    return task_info
        
        return None
    
    def update_task_status(self, task_id: str, status: str, **kwargs) -> bool:
        """
        更新任务状态
        
        Args:
            task_id: 任务ID
            status: 新状态
            **kwargs: 其他更新字段
            
        Returns:
            是否更新成功
        """
        # 在所有用户队列中查找并更新任务
        for user_id, queue in self.user_queues.items():
            for task in queue:
                if task['task_id'] == task_id:
                    task['status'] = status
                    # 更新其他字段
                    for key, value in kwargs.items():
                        task[key] = value
                    return True
        
        return False

    def remove_task(self, task_id: str) -> bool:
        """
        从队列中移除任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            是否移除成功
        """
        # 在所有用户队列中查找并移除任务
        for user_id, queue in self.user_queues.items():
            for i, task in enumerate(queue):
                if task['task_id'] == task_id:
                    queue.pop(i)
                    # 如果用户队列为空，清理相关记录
                    if not queue:
                        del self.user_queues[user_id]
                        if user_id in self.user_last_served:
                            del self.user_last_served[user_id]
                    return True
        
        return False


# 全局实例
user_identifier = UserIdentification()
fair_queue = UserFairQueue(user_identifier)


def get_client_ip(request) -> str:
    """
    从请求中获取客户端IP
    
    Args:
        request: Flask请求对象
        
    Returns:
        客户端IP地址
    """
    # 尝试从各种头部获取真实IP
    if request.headers.get('X-Forwarded-For'):
        return request.headers.get('X-Forwarded-For').split(',')[0].strip()
    elif request.headers.get('X-Real-IP'):
        return request.headers.get('X-Real-IP')
    elif request.headers.get('CF-Connecting-IP'):
        return request.headers.get('CF-Connecting-IP')
    else:
        return request.remote_addr or '127.0.0.1'


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(level=logging.INFO)
    
    # 测试用户识别
    test_ip = "192.168.1.100"
    user_id1 = user_identifier.identify_user(test_ip)
    print(f"用户ID: {user_id1}")
    
    # 测试公平队列
    fair_queue.add_task("task_001", test_ip, {"image": "test.jpg"})
    fair_queue.add_task("task_002", test_ip, {"image": "test2.jpg"})
    fair_queue.add_task("task_003", "192.168.1.101", {"image": "test3.jpg"})
    
    print(f"队列状态: {fair_queue.get_queue_status()}")
    
    # 获取任务
    task = fair_queue.get_next_task()
    print(f"下一个任务: {task}")
